Add surface tracer restoring by katsmith133 · Pull Request #355 · E3SM-Project/Omega

katsmith133 · 2026-03-04T20:31:16Z

This PR adds surface restoring for tracers and updates the tests and docs to reflect these changes.

Passes all CTests on PM-CPU and PM-GPU.

Checklist

Testing

CTest unit tests

Machine: PM-CPU, PM-GPU
Compiler: gnu, gnugpu
Build type: Relase
Result: All tests passed

Tests added/modified/impacted:

CTest: in TendencyTermsTest added test: testSurfaceTracerRestoringOnCell, and in AuxiliaryVarsTest added test: testSurfTracerRestAuxVars

Documentation

Documentation has be built locally here and charges are as expected: https://portal.nersc.gov/project/e3sm/katsmith/omega_docs_surfforce/html/index.html

alicebarthel · 2026-03-04T22:07:39Z

Thanks @katsmith133 ! Were you thinking of adding extra documentation to this? I see that you have updated the existing doc with a mention of these terms. It might be worthwhile adding a description paragraph to document the assumptions we've made on the current version of surface restoring? I don't know what our agreed-upon documentation level is though, since the other terms are not necessarily described either....

components/omega/configs/Default.yml

components/omega/test/ocn/TendencyTermsTest.cpp

components/omega/test/ocn/AuxiliaryVarsTest.cpp

alicebarthel

@katsmith133 this looks good to me!
I concur that all tests pass on pm-gpu.
my main questions/requests are whether we want to document the current surface restoring configuration a little more (e.g. velocity and maxDiff need to be positive, current restoring everywhere, expected T,S variables are the state variables CT, SA).
Do we have a plan to test the minLayerCell different from zero at this point or are we waiting for a later testing configuration?

components/omega/src/ocn/Tendencies.cpp

cbegeman · 2026-03-16T15:06:09Z

@katsmith133 Let me know if this can/should be tested with a single-column polaris test with surface restoring

katsmith133 · 2026-03-17T21:03:17Z

@katsmith133 this looks good to me! I concur that all tests pass on pm-gpu. my main questions/requests are whether we want to document the current surface restoring configuration a little more (e.g. velocity and maxDiff need to be positive, current restoring everywhere, expected T,S variables are the state variables CT, SA). Do we have a plan to test the minLayerCell different from zero at this point or are we waiting for a later testing configuration?

@alicebarthel Thanks for reviewing this! I will add more details to the documentation to cover the things you suggest. Great ideas!

Also, as far as I can tell, the way MPAS-O does this is to apply it to the minLevelCell, not to the zeroth level:
https://github.com/E3SM-Project/E3SM/blob/6918a79638467df967c74c2a544e9a120d5179a7/components/mpas-ocean/src/shared/mpas_ocn_tracer_surface_restoring.F#L143
What kind of testing are you thinking about here? Are you talking about specific to the surface forcing, or just in general when minLayerCell is not zero? If the later, I think that should go in the vertical coordinate testing, not here. If the former, can you give me more detail on what you mean?

katsmith133 · 2026-03-18T16:43:31Z

After a discussion with @alicebarthel, we've noted that this PR does not include the following features:

Flags to run only restoring on certain tracers, for example the ability to only run S restoring and not the others (I will be adding this option to this PR once PM is back from maintenance)
The global sum calculation to ensure a net flux of zero for salinity (do we want to include this in this PR or save this for later dev, just to get this feature in quickly?)
An option to shut off tracer restoring under sea ice (current default in MPAS-O is to restore under sea ice, so this option is inline with that, but maybe we still want the option to turn it off?)

cbegeman · 2026-03-18T19:18:32Z

I'm personally fine with saving (2) and (3) for later. In particular, we'd want to add sea ice forcing fields before (3).

katsmith133 · 2026-03-20T22:03:11Z

Ok, I've added a way to specify which tracers or tracer groups you want to have restored. So you can either specify individual tracers by name in TracersToRestore or whole groups by their tracer group name in TracerGroupsToRestore. Example: in Default.yml under the SurfaceRestoring header you could put:

TracersToRestore: [Temperature, Salinity, Debug1]
or
TracerGroupsToRestore: [Base]

Only one option needs to be used, but if both are used any overlapping tracers (in this case Temp and Salt as they are in the Base tracer group) will just be counted once along with any non-overlapping tracers in both categories. Perhaps this is an overkill way to do this? Would we prefer to just have the TracersToRestore option to make it simpler?

Along with this change I have added a bit of extra testing to TendencyTermsTest.cpp to ensure this capability is covered. So now it tests different combos of tracers being restored (temp only, salt only, and a combo of temp, salt, and one debug).

Will next work on adding more documentation now, just need to figure out where best to add additional information on the surface restoring in the Docs...

One last thought... are we assuming that the restoring files are in the units of SA and CT? Or should we assume they will be in T,S like they are for MPAS-O? I think that might be the better assumption? In that case, some code to convert would be good. I could add a flag that toggles that conversion if we don't want to assume one way or another? Or another option would be to have it decide based upon units in the restoring file? The infrastructure is not there yet to read that file, so this would be a later dev.

katsmith133 · 2026-03-26T20:27:18Z

I've updated the tracer selection to be more simple, added new docs for summarizing the forcing (with @alicebarthel planning to add thermodynamics forcing to this when that is PR'd), and addressed all current reviewer comments. @vanroekel and @mwarusz do you two want to take a look at this? Or should I remove you/ask others.

vanroekel · 2026-03-26T22:08:46Z

I won't be able to review this till early next week. If okay to wait that long I will review monday. If you are ready to go sooner, just remove me.

mwarusz

Thanks for your work on this @katsmith133 It looks mostly good, I have a couple of small comments and two bigger comments regarding performance. If this feature is needed soon, the performance comments can be addressed in a later PR.

Performance comments:

I currently don't see why we need to create an aux array for SurfTracerRestoringDiffsCell. Can't these differences be computed inside the surface restoring tendency term ?

The approach of selecting tracers to restore using a mask seems sub-optimal to me. The way I would have done it is:

On the host, determine the indices of tracers to restore and the total number of tracers to restore

Allocate and fill a device array

  Array1DI4 TracerIdsToRestore("TracerIdsToRestore", NTracersToRestore);

Compute the tendency term as follows:

 parallelFor(
      {NTracersToRestore, Mesh->NCellsAll},
          KOKKOS_LAMBDA(int R, int ICell) {
         const int KMin = MinLayerCell(ICell);
         const int L = TracerIdsToRestore(R);
         LocSurfaceTracerRestoring(LocTracerTend, L, ICell, KMin,
                                      SurfTracerRestoringDiffsCell);
      });

In addition to the suggestions below, if you decide to keep the SurfTracerRestoringDiffsCell aux array, could you add the following piece of code in the auxiliary state test to check that this variable is computed:

   Real SurfTracerRestoringDiffsCSum = 0;
   for (I4 LTracer = 0; LTracer < NTracers; ++LTracer) {
      auto TracerDiffsCell = Kokkos::subview(
          DefAuxState->SurfTracerRestAux.SurfTracerRestoringDiffsCell, LTracer,
          Kokkos::ALL);

      SurfTracerRestoringDiffsCSum += sum(TracerDiffsCell, OnCell);
   }

   if (!Kokkos::isfinite(SurfTracerRestoringDiffsCSum)) {
      Err++;
      LOG_ERROR("AuxStateTest: SurfTracerRestoringDiffsCell FAIL");
   }

You can add it after a similar check here:

Omega/components/omega/test/ocn/AuxiliaryStateTest.cpp

Lines 302 to 308 in c409e9d

    
           const Real Del2TracersCSum = 
        
               sum(DefAuxState->TracerAux.Del2TracersCell, NTracers, NCellsOwned, 
        
                   VCoord->MinLayerCell, VCoord->MaxLayerCell); 
        
           if (!Kokkos::isfinite(Del2TracersCSum)) { 
        
              Err++; 
        
              LOG_ERROR("AuxStateTest: Del2TracersOnCell FAIL"); 
        
           }

components/omega/src/ocn/auxiliaryVars/SurfTracerRestAuxVars.h

components/omega/src/ocn/AuxiliaryState.cpp

components/omega/src/ocn/Tendencies.cpp

components/omega/test/ocn/TendencyTermsTest.cpp

katsmith133 · 2026-03-30T23:06:21Z

@mwarusz Thanks for the review! I think I touched on everything you suggested. Please let me know if my changes addressed your concerns or not.

components/omega/configs/Default.yml

components/omega/src/ocn/Tendencies.cpp

vanroekel

Overall looks good @katsmith133 but I had two bigger questions. I'd also be interested in what @alicebarthel thinks about the maxDiff comment. I've never understood why we included maxDiff = 100 in MPAS-Ocean, it seems so big as to not be useful.

components/omega/doc/devGuide/Forcing.md

cbegeman · 2026-04-02T15:39:54Z

I think it would be pretty approachable to add a surface restoring variant of the single column test E3SM-Project/polaris#435. Let me know if you need help.

katsmith133 · 2026-04-02T15:42:50Z

I think it would be pretty approachable to add a surface restoring variant of the single column test E3SM-Project/polaris#435. Let me know if you need help.

Sure, I think that is a good idea, @cbegeman. I can get started on that later today.

@alicebarthel any thoughts on Luke's MaxDiff comment?

katsmith133 · 2026-04-06T17:23:49Z

Removed MaxDiff code, updated and rebuilt documentation, and addressed all review comments. Passes CTests on PM-CPU and PM-GPU.

As per the conversation this morning, we will save calling in the timestepper and testing in polaris for a different PR once appropriate code is merged.

@alicebarthel, @vanroekel, and @mwarusz, please let me know if you are happy with the current code and we can merge.

Thanks!

…r groups

mwarusz

I have one small suggestion, but otherwise this looks good to me. Passes CTests on frontier-gpu. Thanks for addressing my previous comments @katsmith133

components/omega/src/ocn/Tendencies.cpp

katsmith133 · 2026-04-06T21:30:16Z

Thanks @mwarusz for the review and great suggestions!

alicebarthel · 2026-04-06T21:49:45Z

Thanks for the quick turn-around @katsmith133!
I did a brief visual inspection for the clamp/maxdiff changes and it looks good to me!
Thanks for your work on this, and especially for getting the forcing documentation going!

vanroekel

Looks great from a visual pass. Thanks @katsmith133 for the great work.

mark-petersen · 2026-04-07T20:10:08Z

Merged into develop locally. Passes all cTests with perlmutter gnu CPU and GPU. Passes Polaris manufactured solution test. (that does not test surface forcing, but we didn't mess it up...)

katsmith133 requested review from alicebarthel, mwarusz and vanroekel March 4, 2026 20:31

alicebarthel reviewed Mar 6, 2026

View reviewed changes

components/omega/configs/Default.yml Outdated Show resolved Hide resolved

alicebarthel reviewed Mar 7, 2026

View reviewed changes

components/omega/test/ocn/TendencyTermsTest.cpp Show resolved Hide resolved

alicebarthel reviewed Mar 7, 2026

View reviewed changes

components/omega/test/ocn/AuxiliaryVarsTest.cpp Outdated Show resolved Hide resolved

alicebarthel reviewed Mar 7, 2026

View reviewed changes

alicebarthel reviewed Mar 12, 2026

View reviewed changes

components/omega/src/ocn/Tendencies.cpp Outdated Show resolved Hide resolved

mwarusz reviewed Mar 27, 2026

View reviewed changes

katsmith133 force-pushed the omega/surface-restoring branch from e8bb606 to 625e627 Compare March 30, 2026 22:41

vanroekel reviewed Mar 31, 2026

View reviewed changes

components/omega/configs/Default.yml Outdated Show resolved Hide resolved

vanroekel reviewed Mar 31, 2026

View reviewed changes

components/omega/src/ocn/Tendencies.cpp Show resolved Hide resolved

vanroekel reviewed Mar 31, 2026

View reviewed changes

components/omega/src/ocn/Tendencies.cpp Outdated Show resolved Hide resolved

vanroekel requested changes Mar 31, 2026

View reviewed changes

alicebarthel reviewed Apr 1, 2026

View reviewed changes

components/omega/doc/devGuide/Forcing.md Outdated Show resolved Hide resolved

alicebarthel reviewed Apr 1, 2026

View reviewed changes

components/omega/doc/devGuide/Forcing.md Outdated Show resolved Hide resolved

alicebarthel reviewed Apr 1, 2026

View reviewed changes

components/omega/doc/devGuide/Forcing.md Outdated Show resolved Hide resolved

starts adding surface tracer restoring

dc5a8da

katsmith133 added 12 commits April 6, 2026 11:23

completes working surface forcing

1717bfa

rebase

f142ae5

updates docs

864cf69

adds ability to apply surface restoring to specified tracers or trace…

8777f79

…r groups

removed group tracer optioin in config and checks if maxdiff is pos

4367a5e

adds more documentation and renames vars

4189450

removed layer loop from surface restoring calc

c785491

adds suggestions from pr review

acba12c

rebases and updates docs

56be44b

removes redundent call of OmegaConfig

c6ddfac

reverts submodule changes

4f79be5

removes maxdiff and adds suggestions from pr review

8fa3a02

katsmith133 force-pushed the omega/surface-restoring branch from d016b5b to 8fa3a02 Compare April 6, 2026 19:05

linting

a9e2fa4

mwarusz approved these changes Apr 6, 2026

View reviewed changes

components/omega/src/ocn/Tendencies.cpp Outdated Show resolved Hide resolved

adds suggestion from pr review

761b24d

vanroekel approved these changes Apr 7, 2026

View reviewed changes

mark-petersen self-assigned this Apr 7, 2026

alicebarthel approved these changes Apr 7, 2026

View reviewed changes

mark-petersen merged commit f07cde6 into E3SM-Project:develop Apr 7, 2026
1 check passed

	const Real Del2TracersCSum =
	sum(DefAuxState->TracerAux.Del2TracersCell, NTracers, NCellsOwned,
	VCoord->MinLayerCell, VCoord->MaxLayerCell);
	if (!Kokkos::isfinite(Del2TracersCSum)) {
	Err++;
	LOG_ERROR("AuxStateTest: Del2TracersOnCell FAIL");
	}

Conversation

katsmith133 commented Mar 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Testing

Documentation

Uh oh!

alicebarthel commented Mar 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

alicebarthel left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

cbegeman commented Mar 16, 2026

Uh oh!

katsmith133 commented Mar 17, 2026

Uh oh!

katsmith133 commented Mar 18, 2026

Uh oh!

cbegeman commented Mar 18, 2026

Uh oh!

katsmith133 commented Mar 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

katsmith133 commented Mar 26, 2026

Uh oh!

vanroekel commented Mar 26, 2026

Uh oh!

mwarusz left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

katsmith133 commented Mar 30, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vanroekel left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cbegeman commented Apr 2, 2026

Uh oh!

katsmith133 commented Apr 2, 2026

Uh oh!

katsmith133 commented Apr 6, 2026

Uh oh!

mwarusz left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

katsmith133 commented Apr 6, 2026

Uh oh!

alicebarthel commented Apr 6, 2026

Uh oh!

vanroekel left a comment

Choose a reason for hiding this comment

Uh oh!

mark-petersen commented Apr 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

katsmith133 commented Mar 4, 2026 •

edited

Loading

alicebarthel commented Mar 4, 2026 •

edited

Loading

alicebarthel left a comment •

edited

Loading

katsmith133 commented Mar 20, 2026 •

edited

Loading